from django.db import models
from django.utils import timezone
from django.contrib.postgres.fields import JSONField, ArrayField
import os
from django.conf import settings
from apps.storage.models.dataset import DataSet
from django.contrib.postgres.fields import ArrayField, JSONField
from apps.accounts.models.users import User


class Task(models.Model):
    TYPE = [(0, '训练'), (1, '预测')]
    STATUS = [(0, '等待中'), (1, '运行中'), (2, '完成'), (3, '失败')]
    title = models.CharField(max_length=50, default='')
    pub_data = models.DateTimeField(auto_now_add=True)
    start_time = models.DateTimeField(default=timezone.now)
    end_time = models.DateTimeField(default=timezone.now)
    author = models.CharField(max_length=50, default=None, null=True, blank=True)
    algorithm = models.CharField(max_length=100, default=None, null=True)
    label = models.CharField(max_length=50, default=None, null=True)
    features = ArrayField(models.CharField(max_length=50), default=None, null=True)
    params = JSONField(default=None, blank=True, null=True)
    dataset = models.ForeignKey(DataSet, on_delete=models.CASCADE, null=True, default=None)
    result = JSONField(default=None, blank=True, null=True)
    status = models.IntegerField(choices=STATUS, default=0)
    task_type = models.IntegerField(choices=TYPE, default=0)
    visual_pic = ArrayField(models.CharField(max_length=200), default=None, blank=True, null=True)
    model = models.FileField(upload_to='models/%Y/%m/%d', default=None, blank=True, null=True)
    task_id = models.CharField(max_length=100, default=None, null=True)
    result_file = models.FileField(upload_to='files/%Y/%m/%d', default=None, blank=True, null=True)

    class Meta:
        ordering = ('-pub_data',)

    def get_abs_file_path(self):
        abs_file_path = os.path.join(settings.MEDIA_ROOT, str(self.dataset.file))
        return abs_file_path


class VisualRecognitionTask(models.Model):
    '''
    视觉识别任务，重新写的该类
    '''
    TYPE = [(0, '训练+预测'), (1, '训练'), (2, '预测')]
    STATUS = [(0, '等待中'), (1, '运行中'), (2, '完成'), (3, '失败')]
    title = models.CharField(max_length=50, default='Default')
    description = models.CharField(max_length=254, null=True)
    author = models.ForeignKey(User, on_delete=models.CASCADE, related_name='user_vrtask')
    created_time = models.DateTimeField(auto_now_add=True)
    updated_time = models.DateTimeField(auto_now=True)
    end_time = models.DateTimeField(default=timezone.now)
    task_type = models.IntegerField(choices=TYPE, default=0)
    status = models.IntegerField(choices=STATUS, default=0)
    celery_task_id = models.CharField(max_length=254, default=None, null=True)
    model_dir = JSONField(default=None, blank=True, null=True)
    name_to_label = JSONField(default=None, blank=True, null=True)
    train_result = JSONField(default=None, blank=True, null=True)
    test_result = ArrayField(models.CharField(max_length=254), default=None, blank=True, null=True)

    def to_dict(self):
        '''
        返回字段dict
        :return:
        '''
        d = {
            'id': self.id,
            'vr_task_title': self.title,
            'description': self.description,
            'author': self.author.username,
            'task_type': self.task_type,
            'status': self.status,
            'created_time': self.created_time,
            'updated_time': self.updated_time,
            'end_time': self.end_time,
            'train_result': self.train_result,
            'test_result': self.test_result
        }
        return d

    # def delete(self, *args, **kwargs):
    #     '''
    #     TODO:删除字段的同时，停止其占用的资源
    #     :param args:
    #     :param kwargs:
    #     :return:
    #     '''
    #     pass
